#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图像分支模型
基于ResNet50 backbone，包含特征增强模块
"""

import torch
import torch.nn as nn
import torchvision.models as models
from typing import Tuple


class CAS(nn.Module):
    """
    TODO: Implement custom Channel Attention Module (CAS)
    通道注意力机制模块
    目前使用恒等映射作为占位符
    """
    def __init__(self, in_channels: int = 256):
        super().__init__()
        self.in_channels = in_channels
        # TODO: 实现真实的通道注意力机制

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, 256, H, W) 特征图

        Returns:
            (B, 256, H, W) 增强后的特征图
        """
        # TODO: 实现CAS的真实结构
        return x  # 恒等映射


class CFBConv(nn.Module):
    """
    TODO: Implement custom Cross-scale Feature Convolution (CFBConv)
    跨尺度特征卷积模块
    目前使用恒等映射作为占位符
    """
    def __init__(self, in_channels: int = 256):
        super().__init__()
        self.in_channels = in_channels
        # TODO: 实现真实的跨尺度特征卷积

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, 256, H, W) 特征图

        Returns:
            (B, 256, H, W) 增强后的特征图
        """
        # TODO: 实现CFBConv的真实结构
        return x  # 恒等映射


class SLF(nn.Module):
    """
    TODO: Implement custom Spatial Localization Fusion (SLF)
    空间定位融合模块
    目前使用恒等映射作为占位符
    """
    def __init__(self, in_channels: int = 256):
        super().__init__()
        self.in_channels = in_channels
        # TODO: 实现真实的空间定位融合机制

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, 256, H, W) 特征图

        Returns:
            (B, 256, H, W) 增强后的特征图
        """
        # TODO: 实现SLF的真实结构
        return x  # 恒等映射


class ImageBranch(nn.Module):
    """
    图像分支模型
    基于ResNet50 backbone，输出256维特征图
    """

    def __init__(self, pretrained: bool = True):
        super().__init__()

        # 加载预训练的ResNet50 backbone
        self.backbone = models.resnet50(pretrained=pretrained)

        # 移除最后的平均池化和全连接层
        self.backbone = nn.Sequential(*list(self.backbone.children())[:-2])

        # 添加1x1卷积将通道数调整为256
        self.channel_adapter = nn.Conv2d(2048, 256, kernel_size=1, stride=1, padding=0)
        self.bn_adapter = nn.BatchNorm2d(256)
        self.relu = nn.ReLU(inplace=True)

        # 特征增强模块
        self.cas = CAS(in_channels=256)
        self.cfb = CFBConv(in_channels=256)
        self.slf = SLF(in_channels=256)

        # 输出层
        self.output_bn = nn.BatchNorm2d(256)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播

        Args:
            x: (B, 3, 224, 224) 输入图像

        Returns:
            (B, 256, H, W) 256维特征图
        """
        # 通过ResNet50 backbone
        features = self.backbone(x)  # (B, 2048, H, W)

        # 通道调整
        features = self.channel_adapter(features)  # (B, 256, H, W)
        features = self.bn_adapter(features)
        features = self.relu(features)

        # 特征增强（依次通过三个模块）
        features = self.cas(features)   # 通道注意力增强
        features = self.cfb(features)   # 跨尺度特征卷积
        features = self.slf(features)   # 空间定位融合

        # 输出批归一化
        features = self.output_bn(features)

        return features  # (B, 256, H, W)


def test_image_branch():
    """测试图像分支模型"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 创建模型
    model = ImageBranch(pretrained=False).to(device)

    # 创建测试输入
    batch_size = 4
    test_input = torch.randn(batch_size, 3, 224, 224).to(device)

    # 前向传播
    with torch.no_grad():
        output = model(test_input)
        print(f"输入尺寸: {test_input.shape}")
        print(f"输出尺寸: {output.shape}")
        print(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")

    return output.shape


if __name__ == "__main__":
    test_image_branch()